import os
import sys
import math
import numpy as np
import matplotlib.pyplot as plt

def cal_distance(speed1, direct1, speed2, direct2):
    # return (speed1 - speed2) ** 2
    u1 = speed1 * math.sin(direct1)
    v1 = speed1 * math.cos(direct1)
    u2 = speed2 * math.sin(direct2)
    v2 = speed2 * math.cos(direct2)
    return ((u1 - u2) ** 2 + (v1 - v2) ** 2) ** 0.5


def cal_distance2(speed1, speed2):
    return (speed1 - speed2) ** 2



def normalization(data):
    avg = x = np.abs(data).mean(axis=0)
    data /= avg
    return data, avg



def denormalization(data, max1, min1):
    data *= (max1 - min1)
    data += min1
    return data



def cal(data1, start1, end1, data2, start2, end2):
    """
    计算data1[start1: end1)和data2[start2: end2)的最小差距，并记录匹配规则

    return:
        min_v: 最小差距
        rule: 匹配规则
            rule中包含若干[a, b]二元组，表示对应序号的点应该匹配
    """
    dp = [[sys.float_info.max for _ in range(end2 - start2)] for _ in range(end1 - start1)]
    dp[0][0] = cal_distance(data1[start1, 0], data2[start2, 0])
    for i in range(end1 - start1 - 1):
        for j in range(end2 - start2 - 1):
            dp[i + 1][j] = min(dp[i + 1][j], dp[i][j] + cal_distance(data1[start1 + i + 1, 0], data2[start2 + j, 0]))
            dp[i][j + 1] = min(dp[i][j + 1], dp[i][j] + cal_distance(data1[start1 + i, 0], data2[start2 + j + 1, 0]))
            dp[i + 1][j + 1] = min(dp[i + 1][j + 1], dp[i][j] + cal_distance(data1[start1 + i + 1, 0], data2[start2 + j + 1, 0]))
    
    cur1 = end1 - start1 - 1
    cur2 = end2 - start2 - 1
    rule = [[cur1, cur2]]
    min_v = dp[cur1][cur2]
    while cur1 != 0 or cur2 != 0:
        lastv = dp[cur1][cur2]

        if cur1 > 0 and dp[cur1 - 1][cur2] <= lastv:
            lastv = dp[cur1 - 1][cur2]
            last1 = cur1 - 1
            last2 = cur2

        if cur2 > 0 and dp[cur1][cur2 - 1] <= lastv:
            lastv = dp[cur1][cur2 - 1]
            last1 = cur1
            last2 = cur2 - 1

        if cur1 > 0 and cur2 > 0 and dp[cur1 - 1][cur2 - 1] <= lastv:
            lastv = dp[cur1 - 1][cur2 - 1]
            last1 = cur1 - 1
            last2 = cur2 - 1


        cur1 = last1
        cur2 = last2
        rule.append([last1, last2])

    return min_v, rule



def cal2(data1, start1, data2, start2, seq1, min_len2, max_len2, limit):
    """
    给出data1和start1，data2和start2，找到长度为[min_len2, max_len2]之间平均匹配距离最小的上游匹配序列，并记录匹配规则

    return:
        end2：最佳匹配区间终点
        min_v: 最小平均匹配距离
        rule: 匹配规则
            rule中包含若干[a, b]二元组，表示对应序号的点应该匹配
    """
    if seq1 <= 1:
        return None

    dp = [[sys.float_info.max for _ in range(max_len2)] for _ in range(seq1)]
    last = [[[-1, -1] for _ in range(max_len2)] for _ in range(seq1)]
    dp[0][0] = cal_distance2(data1[start1, 0], data2[start2, 0])

    compress_count_i = [0 for _ in range(len(dp))]
    compress_count_j = [0 for _ in range(len(dp[0]))]
    
    for i in range(len(dp) - 1):
        for j in range(len(dp[0]) - 1):
            v10 = dp[i][j] + cal_distance2(data1[start1 + i + 1, 0].item(), data2[start2 + j, 0].item())
            if dp[i + 1][j] > v10 and compress_count_i[i + 1] < limit and compress_count_j[j] < limit:
                dp[i + 1][j] = v10
                compress_count_i[i + 1] += 1
                compress_count_j[j] += 1
                if last[i + 1][j] != [-1, -1]:
                    x, y = last[i + 1][j]
                    compress_count_i[x] -= 1
                    compress_count_j[y] -= 1
                last[i + 1][j] = [i, j]
            
            v01 = dp[i][j] + cal_distance2(data1[start1 + i, 0].item(), data2[start2 + j + 1, 0].item())
            if dp[i][j + 1] > v01 and compress_count_i[i] < limit and compress_count_j[j + 1] < limit:
                dp[i][j + 1] = v01
                compress_count_i[i] += 1
                compress_count_j[j + 1] += 1
                if last[i][j + 1] != [-1, -1]:
                    x, y = last[i][j + 1]
                    compress_count_i[x] -= 1
                    compress_count_j[y] -= 1
                last[i][j + 1] = [i, j]
            
            v11 = dp[i][j] + cal_distance2(data1[start1 + i + 1, 0].item(), data2[start2 + j + 1, 0].item())
            if dp[i + 1][j + 1] > v11 and compress_count_i[i + 1] < limit and compress_count_j[j + 1] < limit:
                dp[i + 1][j + 1] = v11
                compress_count_i[i + 1] += 1
                compress_count_j[j + 1] += 1
                if last[i + 1][j + 1] != [-1, -1]:
                    x, y = last[i + 1][j + 1]
                    compress_count_i[x] -= 1
                    compress_count_j[y] -= 1
                last[i + 1][j + 1] = [i, j]

    cur1 = seq1 - 1
    cur2 = -1
    min_v = sys.float_info.max
    for i in range(min_len2, max_len2):
        value = dp[cur1][i]
        if min_v > value:
            min_v = value
            cur2 = i
    
    end2 = cur2 + start2 + 1
    rule = [[cur1, cur2]]
    while last[cur1][cur2] != [-1, -1]:
        rule.append(last[cur1][cur2])
        cur1, cur2 = last[cur1][cur2]
    # print_to_file(dp, start1, start1 + seq, start2, end2, min_v, rule)
    return end2, min_v, rule



def print_to_file(dp, start1, end1, start2, end2, min_v, rule):
    folder = "dp_record"
    file = f"dp_{start1}_{end1}_{start2}_{end2}_{min_v}"
    if not os.path.exists(folder):
        os.mkdir(folder)
    path = os.path.join(folder, file)
    f = open(path, "w")
    fout = sys.stdout
    sys.stdout = f
    print("dp")
    for i in range(len(dp)):
        for j in range(len(dp[i])):
            print("{:04.4f}".format(dp[i][j]), end=" ")
        print()
    print()
    print(f"min_v:{min_v}")
    print("rule")
    for x, y in rule:
        print(x, y)
    sys.stdout = fout


def draw(origin_data1, origin_data2, data_name, dataset_start, dataset_end, features, res):
    folder_path = os.path.join(data_name, f"test_{dataset_start}_{dataset_end}")
    if not os.path.exists(folder_path):
        os.mkdir(folder_path)


    for feature in features:
        for start1, end1, start2, end2, rule in res:
            target_data = origin_data1[feature]
            source_data = origin_data2[feature]
            y_down = []
            y_up = []
            
            """
            for i in range(-1, -len(rule) - 1, -1):
                p1 = rule[i][0]
                p2 = rule[i][1]
                y_down.append(target_data[p1])
                y_up.append(source_data[p2])
            """

            for y in target_data[start1: end1]:
                y_down.append(y)
            for y in source_data[start2: end2]:
                y_up.append(y)
            x = range(len(y_up))
            plt.plot(x, y_up, color='red')
            x = range(len(y_down))
            plt.plot(x, y_down, color='blue')

            x = range(max(len(y_down), len(y_up)))
            for p in rule:
                p1 = p[0]
                p2 = p[1]
                x_connect = [x[p1], x[p2]]
                y_connect = [y_down[p1], y_up[p2]]
                plt.plot(x_connect, y_connect, linestyle='--', color='gray')
            
            plt.savefig(os.path.join(folder_path, f"{feature}_{start1}_{end1}_{start2}_{end2}.png"))
            plt.cla()